# Basic code for cutoff function.

import jax
import jax.numpy as jnp
import numpy as np
import haiku as hk

from typing import Optional, Type, Union, List, Tuple

_CUTOFF_BY_KEY = dict()
def _cutoff_register(*aliases):
    """Return the alias register."""
    def alias_reg(cls):
        name = cls.__name__
        name = name.lower()
        if name not in _CUTOFF_BY_KEY:
            _CUTOFF_BY_KEY[name] = cls

        for alias in aliases:
            if alias not in _CUTOFF_BY_KEY:
                _CUTOFF_BY_KEY[alias] = cls

        return cls

    return alias_reg

class Cutoff(hk.Module):
    r"""Cutoff function.

    Args:
        cutoff (float):   Cutoff distance.
    
    """

    def __init__(self,
                 fp_type = jnp.float32, 
                 cutoff: Optional[float] = None,
                 name: str = "cutoff",
                 ):
        super().__init__(name=name)

        self.fp_type = fp_type

        self.cutoff = jnp.asarray(cutoff, dtype=self.fp_type)

    def __call__(self,
                 distance: jnp.ndarray,
                 mask: Optional[jnp.ndarray] = None,
                 cutoff: Optional[jnp.ndarray] = None,
                 ) -> Tuple[jnp.ndarray, jnp.ndarray]:
        r"""Compute cutoff.
        
        Args:
            distance (Distance):    Array of shape (A, A). Distance between atoms.
            mask (Mask):            Array of shape (A, A). Mask for distance.
            cutoff (Cutoff):        Array of shape (A, A). Cutoff distance. Default: None.
        
        Returns:
            decay (Array): Array of shape (A, A). Data type is float.
            mask (Array):  Array of shape (A, A). Data type is bool.
        """

        raise NotImplementedError